use std::sync::Arc;

use chatglm_cpp_sys::Pipeline;
use talktive::{Message, TalkTive, TalktiveKenerl, TalktiveKenerlAns};

use crate::{config::ChatGLMCppConfig, history, role::Role, tool_sig};

pub struct ChatGLMCpp {
    pipeline: Arc<Pipeline>,
}

unsafe impl Sync for ChatGLMCpp {}
impl ChatGLMCpp {
    fn decode_res(res : &str) -> anyhow::Result<TalktiveKenerlAns<Role>> {
        if res.ends_with("<|observation|>") {
            todo!()
        } else {
            let mut res = res.to_owned();
            res.remove_matches("<|assistant|>");
            Ok(TalktiveKenerlAns::Message(Message::from_parts(Role::Assistant, res)))
        }
    }
}

//#[async_trait::async_trait]
impl TalktiveKenerl for ChatGLMCpp {
    type KenerlConfig = ChatGLMCppConfig;

    type ToolSigBatchEncodedType = String;

    type HistoryBatchEncodedType = String;

    type ToolSigEncoder = tool_sig::Encoder;

    type HistoryEncoder = history::Encoder;

    fn name() -> &'static str {
        "ChatGLMCpp"
    }

    fn build(config: &Self::KenerlConfig) -> anyhow::Result<Self>
    where
        Self: Sized,
    {
        let pipeline = Pipeline::new(&config.model_path, &config.generate_config)?;
        Ok(Self {
            pipeline: Arc::new(pipeline),
        })
    }

    async fn send(
        &self,
        history: Option<Self::HistoryBatchEncodedType>,
        tools: Option<Arc<Self::ToolSigBatchEncodedType>>,
    ) -> anyhow::Result<
        Vec<
            talktive::TalktiveKenerlAns<
                <Self::HistoryEncoder as talktive::HistoryEncoder>::MessageRole,
            >,
        >,
    > {
        let mut prompt = match history {
            Some(s) => s,
            None => String::new(),
        };
        if let Some(tools) = tools {
            prompt.push_str(&tools);
        }
        let res = self.pipeline.generate(&prompt);
        let res = Self::decode_res(res.trim())?;
        Ok(vec![res])
    }

    async fn stop(&self) -> anyhow::Result<()> {
        Ok(())
    }
}
